package ccnl.demo;

import org.apache.spark.Accumulator;
import org.apache.spark.SparkContext;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.api.java.function.DoubleFlatMapFunction;
import org.apache.spark.api.java.function.Function;
import org.apache.spark.mllib.linalg.Vector;
import org.apache.spark.mllib.regression.LabeledPoint;
import org.apache.spark.sql.DataFrame;
import org.apache.spark.sql.Row;
import org.apache.spark.util.StatCounter;
import org.codehaus.janino.Java;
import org.dmg.pmml.False;
import scala.Tuple2;

/**
 * Created by wong on 16/4/3.
 */
public class JrddTools {
    public static long getPosNums(JavaRDD<LabeledPoint> jrdd1) {
        JavaRDD<LabeledPoint> jrdd2 = jrdd1.filter(v1 -> v1.label() == 1);
        return jrdd2.count();
    }
    public static long getNegNums(JavaRDD<LabeledPoint> jrdd1){
        JavaRDD<LabeledPoint> jrdd2 = jrdd1.filter(v1 -> v1.label() == 0);
        return jrdd2.count();
    }
    public static JavaRDD<LabeledPoint> convertRow2Lp(JavaRDD<Row> jrdd_row) {
        JavaRDD<LabeledPoint> jrdd_lp = jrdd_row.map(v1 -> {
            final int index = v1.fieldIndex("features");
            int label_index = v1.fieldIndex("C1");
            LabeledPoint lp = new LabeledPoint(Double.valueOf(v1.getString(label_index)), (Vector) v1.get(index));
            return lp;
        });
        return jrdd_lp;
    }
    public static JavaRDD<Row> getTrain(JavaRDD<Row> jrdd) {
        JavaRDD<Row> jrdd_train = jrdd.filter(v1 -> {
            int index = v1.fieldIndex("C2");
            String s = v1.getString(index);
            return !s.substring(8, 10).equals("21");
        });
        return jrdd_train;
    }
    public static JavaRDD<Row> getTest(JavaRDD<Row> jrdd){
        JavaRDD<Row> jrdd_test=jrdd.filter(v1 -> {
            int index=v1.fieldIndex("C2");
            String s=v1.getString(index);
            return s.substring(8, 10).equals("21");
        });
        return jrdd_test;
    }
    public static JavaRDD<Row> getSampledTrain(JavaRDD<Row> train, double frac) {
        JavaRDD<Row> pos_in_train = train.filter(v1 -> {
            int index = v1.fieldIndex("C1");
            return v1.getInt(index) == 1;
        });
        JavaRDD<Row> neg_in_train = train.filter(v1 -> {
            int index = v1.fieldIndex("C1");
            return v1.getInt(index) == 0;
        });
        //JavaRDD<Row> useful_neg=neg_in_train.sample(false,0.00625);
        double truefrac = frac * (double) pos_in_train.count() / neg_in_train.count();
        JavaRDD<Row> useful_neg = neg_in_train.sample(false, truefrac);
        JavaRDD<Row> usefulTrain = pos_in_train.union(useful_neg);

        return usefulTrain;
    }

}
